// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <kllm/core/task_slot.h>
#include <kllm/core/queue.h>

namespace kllm {


    struct server_response {
        // for keeping track of all tasks waiting for the result
        std::unordered_set<int> waiting_task_ids;

        // the main result queue
        std::vector<TaskResult> queue_results;

        std::mutex mutex_results;
        std::condition_variable condition_results;

        // add the id_task to the list of tasks waiting for response
        void add_waiting_task_id(int id_task) {
            SRV_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting_task_ids.size());

            std::unique_lock<std::mutex> lock(mutex_results);
            waiting_task_ids.insert(id_task);
        }

        void add_waiting_tasks(const std::vector<server_task> & tasks) {
            std::unique_lock<std::mutex> lock(mutex_results);

            for (const auto & task : tasks) {
                SRV_DBG("add task %d to waiting list. current waiting = %d (before add)\n", task.id, (int) waiting_task_ids.size());
                waiting_task_ids.insert(task.id);
            }
        }

        // when the request is finished, we can remove task associated with it
        void remove_waiting_task_id(int id_task) {
            SRV_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting_task_ids.size());

            std::unique_lock<std::mutex> lock(mutex_results);
            waiting_task_ids.erase(id_task);
        }

        void remove_waiting_task_ids(const std::unordered_set<int> & id_tasks) {
            std::unique_lock<std::mutex> lock(mutex_results);

            for (const auto & id_task : id_tasks) {
                SRV_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting_task_ids.size());
                waiting_task_ids.erase(id_task);
            }
        }

        // This function blocks the thread until there is a response for one of the id_tasks
        TaskResult recv(const std::unordered_set<int> & id_tasks) {
            while (true) {
                std::unique_lock<std::mutex> lock(mutex_results);
                condition_results.wait(lock, [&]{
                    return !queue_results.empty();
                });

                for (int i = 0; i < (int) queue_results.size(); i++) {
                    if (id_tasks.find(queue_results[i].id) != id_tasks.end()) {
                        TaskResult res = queue_results[i];
                        queue_results.erase(queue_results.begin() + i);
                        return res;
                    }
                }
            }

            // should never reach here
        }

        // single-task version of recv()
        TaskResult recv(int id_task) {
            std::unordered_set<int> id_tasks = {id_task};
            return recv(id_tasks);
        }

        // Send a new result to a waiting id_task
        void send(TaskResult & result) {
            SRV_DBG("sending result for task id = %d\n", result.id);

            std::unique_lock<std::mutex> lock(mutex_results);
            for (const auto & id_task : waiting_task_ids) {
                if (result.id == id_task) {
                    SRV_DBG("task id = %d moved to result queue\n", result.id);

                    queue_results.push_back(std::move(result));
                    condition_results.notify_all();
                    return;
                }
            }
        }
    };

    static nlohmann::ordered_json format_detokenized_response(const std::string &content) {
        return nlohmann::ordered_json{
                {"content", content}
        };
    }

}  // namespace kllm

